'''
Construct tree programs in Python.
The trees are made up of nodes that have some number of child nodes.
Some of the nodes will return parameters passed to the program
some will return constants
some will return operations on their child nodes
'''

from random import random, randint, choice
from copy import deepcopy
from math import log

class fwrapper:
	'''
	a wrapper for the functions that will be used on function nodes.
	its member variables are the name of the function, the function itself, and the number of parameters it takes
	'''
	def __init__(self, function, childcount, name):
		self.function = function
		self.childcount = childcount
		self.name = name
	

class node:
	'''
	the class for function nodes.
	Initialized with a fwrapper. When evluate is called, it evaluates the child nodes and then applies the function to their results
	'''
	def __init__(self, fw, children):
		self.function = fw.function
		self.name = fw.name
		self.children = children
	
	def evaluate(self, inp):
		results = [n.evaluate(inp) for n in self.children]
		return self.function(results)
	
	def display(self, indent = 0):
		print (' '*indent)+self.name
		for c in self.children:
			c.display(indent + 1)


class paramnode:
	'''
	The class for nodes that only return one of the parameters passed to the program.
	Its evaluate method returns the parameter specified by idx
	'''
	def __init__(self, idx):
		self.idx = idx
	
	def evaluate(self, inp):
		return inp[self.idx]
		
	def display(self, indent = 0):
		print '%sp%d' % (' '*indent, self.idx)

class constnode:
	'''
	Nodes that returns a constant value.
	The evaluate method simply returns teh value with which it is initialized.
	'''
	def __init__(self, v):
		self.v = v
	
	def evaluate(self, inp):
		return self.v
	
	def display(self, indent = 0):
		print '%s%d' % (' '*indent, self.v)


addw = fwrapper(lambda l:l[0]+l[1], 2, 'add')
subw = fwrapper(lambda l:l[0]-l[1], 2, 'subtract')
mulw = fwrapper(lambda l:l[0]*l[1], 2, 'multiply')

def iffunc(l):
	if l[0]>0: 
		return l[1]
	else:
		return l[2]

ifw = fwrapper(iffunc, 3, 'if')

def isgreater(l):
	if l[0]>l[1]:
		return 1
	else:
		return 0

gtw = fwrapper(isgreater, 2, 'isgreater')

flist = [addw, mulw, ifw, gtw, subw]


def exampletree():
	return node(ifw, [node(gtw,[paramnode(0), constnode(3)]), node(addw, [paramnode(1), constnode(5)]), node(subw, [paramnode(1), constnode(2)])])

'''
Creating the initial population
Creating a random program consists of:
creating a root node with a random associated function and then creating as many random child nodes as necessary, which 
in turn may have their own associated random child nodes.

This is most easily defined recursively.
'''

def makerandomtree(pc, maxdepth = 4, fpr = 0.5, ppr = 0.6):
	'''
	pc: number of parameters that the tree will take as input
	fpr: probability that the new node created will be a function node
	ppr: probability that it will be a paramnode if it is not a function node
	'''
	if random() < fpr and maxdepth > 0:
		f = choice(flist)
		children = [makerandomtree(pc, maxdepth-1, fpr, ppr) for i in range(f.childcount)]
		return node(f, children)
	elif random() < ppr:
		return paramnode(randint(0, pc-1))
	else:
		return constnode(randint(0,10))


def scorefunction(tree, s):
	'''
	measuring how good a solution is.
	'''
	dif = 0.0
	for data in s:
		v = tree.evaluate([data[0], data[1]])
		dif += (abs(v-data[2]))/data[2]
	dif = dif/len(s)
	return dif

'''
Mutating programs
Mutation takes a single program and alters it slightly.
The tree program can be altered in a number of ways, by changing the function on a node or by altering its branches.
'''

def mutate(t, pc, probchange = 0.1):
	if random() < probchange:
		return makerandomtree(pc)
	else:
		result = deepcopy(t)
		if isinstance(t, node):
			result.children = [mutate(c, pc, probchange) for c in t.children]
		return result


'''
Crossover
The other type of program modification is crossover or breeding.
Taking two successful programs and combining them to create a new program, usually by replacing a branch from one with a branch from another.
'''

def crossover(t1, t2, probswap = 0.7, top = 1):
	if random() < probswap and not top:
		return deepcopy(t2)
	else:
		result = deepcopy(t1)
		if isinstance(t1, node) and isinstance(t2, node):
			result.children = [crossover(c, choice(t2.children), probswap, 0) for c in t1.children]
		return result


'''
Evolve
'''

def evolve(pc, popsize, rankfunction, maxgen = 500, mutationrate = 0.1, breedingrate = 0.4, pexp = 0.7, pnew = 0.05, precision = 0.05):
	
	def selectindex():
		return int(log(random())/log(pexp))
	
	population =[makerandomtree(pc) for i in range(popsize)]
	for i in range(maxgen):
		scores = rankfunction(population)
		#print scores[0][0]
		if scores[0][0]<precision:
			break
		
		newpop = [scores[0][1], scores[1][1]]
		
		while len(newpop) < popsize:
			if random() > pnew:
				newpop.append(mutate(crossover(scores[selectindex()][1], scores[selectindex()][1], probswap = breedingrate), pc, probchange = mutationrate))
			else:
				newpop.append(makerandomtree(pc))
		population = newpop
	print scores[0][0]
	scores[0][1].display()
	return scores[0][1]


def getrankfunction(dataset):
	def rankfunction(population):
		scores = [(scorefunction(t, dataset), t) for t in population]
		scores.sort()
		return scores
	return rankfunction


def hiddenfunction(x, y):
	return x**2.0 + 2.0*y + 3.0*x + 5.0

def buildhiddenset():
	rows = []
	for i in range(200):
		x = randint(0,40)
		y = randint(0,40)
		rows.append([float(x), float(y), hiddenfunction(x, y)])
	print rows
	return rows

